import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
import joblib

# 1. 数据读取和预处理
def load_data(file_path):
    data = pd.read_csv(file_path)
    X = data[['Voltage']]  # 特征列
    y = data['Distance']   # 目标列
    return X, y

# 2. 模型训练
def train_model(X, y):
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    model = LinearRegression()
    model.fit(X_train, y_train)
    return model

# 3. 保存模型
def save_model(model, file_path):
    joblib.dump(model, file_path)

# 主程序
if __name__ == "__main__":
    # 数据文件路径
    data_file_path = "hall_effect_data.csv"
    model_file_path = "hall_effect_model.pkl"

    # 加载数据
    X, y = load_data(data_file_path)

    # 训练模型
    model = train_model(X, y)

    # 保存模型
    save_model(model, model_file_path)
    print("模型已保存到", model_file_path)
